library(nomep)

### Takes a plist post, a gamma object representing various types of the
### dv (rows are individuals to model, columns are types; diag(4)
### generates four individuals of pure type matrix(c(0, 1, 1, 1),
### nrow=1) would be a single individual belonging to three of a four
### possible classes, etc), the name of a variable to plot (ie a column
### name in post$y), names of the individuals in gamma for the plot
### legend, a hash describing the case type to model, a 2-vector of
### upper and lower bounds functions for within-sample values of
### "variable" to plot across (defaults to c(max, min)), and the number
### of values to plot within those bound.  By default, the function
### holds all variables except "variable" at their means but the case
### hash allows you to relax this.  If non-null, the case hash should
### map variable names (ie columns in post$y) to single-argument
### functions that return a scalar value.  For example, the argument
### hash("ingov" = median) tells the plotting function to hold the
### variable "ingov" at its median, and all other variables at their
### means.  The point.func argument sets the function to use to generate
### point estimates.  Should be a 1-arg function that returns a scalar.
###
### Also takes a variety of plotting details, including point estimate
### color, line type and line weight, whether or not to plot HPD
### intervals, the probability for the HPD intervals, the color, line
### type, and weight of the HPD line segments, and a thining value for
### the HPD lines (instead of printing a line with every plotted point,
### put down a HPD line for every hpd.thin plotted point.  The function
### slightly jitters HPD interval lines across predictions so that they
### don't overlap exactly.  The legend values are pretty self-explanatory.
### Additional arguments are passed on to plot.
###
### Finally, setting return.probs to true causes the function to return the
### full predicted probability matrix.
  
plot.plist.postpred.contin <- function(post, gamma, variable,
                                       gamma.names, case.hash=hash(),
                                       lu=c(min, max), rows=100,
                                       point.func=mean,
                                       point.col="black",
                                       point.lty=1:nrow(gamma),
                                       point.lwd=1, show.hpd=TRUE,
                                       hpd.prob=0.95, hpd.col="grey",
                                       hpd.lty=1, hpd.lwd=0.5,
                                       hpd.thin=2, hpd.jitter=TRUE,
                                       axis.labels=NULL,
                                       show.legend=TRUE,
                                       legend.x=0, legend.y=1,
                                       legend.horiz=TRUE,
                                       legend.box.lwd=0,
                                       legend.cex=1,
                                       legend.width=NULL,
                                       legend.xpd=FALSE,
                                       legend.bty='o',
                                       return.probs=FALSE, ...)
{
    ## Build the hypothetical data matrix, using mean when the case hash
    ## doesn't return a function.  This gives us a ncol(post$y) vector
    ## of fixed values for each variable/column in post$y
    y.case = apply(post$y, 2, mean)
    for (i in 1:ncol(post$y)) {
        func = case.hash[[colnames(post$y)[i]]]
        if (!is.null(func))
            y.case[i] <- func(post$y[,i])
    }

    ## Now expand the vector into a rows-length matrix  
    y <- matrix(rep(y.case, rows), nrow=rows, ncol=length(y.case), byrow=T)
    colnames(y) <- colnames(post$y)
    
    ## Sequence the variable of interest from the lower to upper bound, with
    ## rows levels.  In other words, holding every other variable constant,
    ## give us smooth variation in "variable".  We apply the two functions in
    ## the lu vector to generate the min and max values of this sequence.
    y[,variable] <- seq(lu[[1]](post$y[,variable]),
                        lu[[2]](post$y[,variable]), length.out=rows)

  
    ## Calculate predictions
    p <- predprob.plist.fit(post, gamma, y, hpd.prob, point.func)
    plot(p[[1]][,3], ylim=c(0, 1), type='n', axes=F,
         ylab="Predicted Probability of Selection", ...)

    ## Set up plotting style vectors
    point.col = rep(point.col, length.out=nrow(gamma))
    point.lty = rep(point.lty, length.out=nrow(gamma))
    point.lwd = rep(point.lwd, length.out=nrow(gamma))
    show.hpd = rep(show.hpd, length.out=nrow(gamma))
    hpd.col = rep(hpd.col, length.out=nrow(gamma))
    hpd.lty = rep(hpd.lty, length.out=nrow(gamma))
    hpd.lwd = rep(hpd.lwd, length.out=nrow(gamma))

    ## Plot the confidence intervals, offsetting them a bit for readability
    hpd.offset <- seq(0, hpd.thin - hpd.thin/2, length.out=nrow(gamma))
    if (!hpd.jitter)
        hpd.offset <- rep(0, nrow(gamma))
    print(hpd.offset)
    #hpd.i <- 1:rows
    #hpd.i <- hpd.i[1:rows %% hpd.thin == 0]
    hpd.i <- 1:rows
    for (i in 1:nrow(gamma))
        if (show.hpd[i]) {
          polygon(c(hpd.i, rev(hpd.i)), c(p[[i]][hpd.i,1], rev(p[[i]][hpd.i,3])),
                  density=NA, col=hpd.col[i], border=NA)
        }
          #segments(hpd.i + hpd.offset[i],
          #           p[[i]][hpd.i,1],
          #           hpd.i + hpd.offset[i],
          #           p[[i]][hpd.i,3],
          #           col=hpd.col[i], lty=hpd.lty[i], lwd=hpd.lwd[i])

    for (i in 1:nrow(gamma))
        points(p[[i]][,2], type='l', lty=point.lty[i], lwd=point.lwd[i],
               col=point.col[i])

    if (show.legend)
        legend(legend.x, legend.y, gamma.names, col=point.col,
               lty=point.lty, lwd=point.lwd, horiz=legend.horiz,
               box.lwd=legend.box.lwd, cex=legend.cex,
               text.width=legend.width, xpd=legend.xpd, bty=legend.bty)
    axis(2)
    if (is.null(axis.labels))
        axis(1, at=seq(0, rows, length.out=6),
             labels=round(c(y[1,variable],
                 y[seq(0, rows, length.out=6)[-1], variable]), 2))
    else if (length(axis.labels) == 2)
        axis(1, at=c(1,2), labels=axis.labels)
    else
        axis(1, at=seq(0, rows,
                       length.out=length(axis.labels)),
             labels=axis.labels)
        
    print(seq(lu[[1]](post$y[,variable]),
                       lu[[2]](post$y[,variable]),
                       length.out=length(axis.labels)))
    
    if (return.probs)
       p <- predprob.plist.fit(post, gamma, y, return.summary=FALSE)
}

plot.plist.postpred.dummy <- function(post, gamma, variable,
                                      gamma.names, case.hash=hash(),
                                      point.func=mean,
                                      point.col="black",
                                      point.cex=1,
                                      point.pch=1:nrow(gamma),
                                      point.lwd=1,
                                      show.hpd=TRUE,
                                      hpd.prob=0.95, hpd.col="grey",
                                      hpd.lty=1, hpd.lwd=0.5,
                                      axis.labels=NULL,
                                      show.legend=TRUE,
                                      legend.x=0, legend.y=1,
                                      legend.horiz=TRUE,
                                      legend.box.lwd=0,
                                      legend.cex=1,
                                      legend.width=NULL,
                                      legend.xpd=FALSE,
                                      legend.bty='o',
                                      return.probs=FALSE, ...)
{
    ## Build the hypothetical data matrix, using mean when the case hash
    ## doesn't return a function.  This gives us a ncol(post$y) vector
    ## of fixed values for each variable/column in post$y
    y.case = apply(post$y, 2, mean)
    for (i in 1:ncol(post$y)) {
        func = case.hash[[colnames(post$y)[i]]]
        if (!is.null(func))
            y.case[i] <- func(post$y[,i])
    }

    ## Now expand the vector into a 2-row matrix  
    y <- matrix(rep(y.case, 2), nrow=2, ncol=length(y.case), byrow=T)
    colnames(y) <- colnames(post$y)
    
    ## Sequence the variable of interest from the lower to upper bound, with
    ## rows levels.  This is a dummy, so that is 0 and 1.
    y[,variable] <- 0:1

  
    ## Calculate predictions
    p <- predprob.plist.fit(post, gamma, y, hpd.prob, point.func)
    plot(p[[1]][,3], ylim=c(0, 1), xlim=c(-0.5, 1.5), type='n', axes=F,
         ylab="Predicted Probability of Selection", ...)

    ## Set up plotting style vectors
    point.col = rep(point.col, length.out=nrow(gamma))
    point.pch = rep(point.pch, length.out=nrow(gamma))
    point.cex = rep(point.cex, length.out=nrow(gamma))
    point.lwd = rep(point.lwd, length.out=nrow(gamma))
    show.hpd = rep(show.hpd, length.out=nrow(gamma))
    hpd.col = rep(hpd.col, length.out=nrow(gamma))
    hpd.lty = rep(hpd.lty, length.out=nrow(gamma))
    hpd.lwd = rep(hpd.lwd, length.out=nrow(gamma))


    ## Setup the x locations.  We put the 0 values between -.25 and
    ## .25 and the 1 values between .75 and 1.25
    xvals <- matrix(NA, nrow=2, ncol=nrow(gamma))
    xvals[1,] <- seq(-.25, .25, length.out=nrow(gamma))
    xvals[2,] <- seq(.75, 1.25, length.out=nrow(gamma))

    for (i in 1:nrow(gamma)) {
        if (show.hpd[i])
            segments(xvals[,i], p[[i]][,1], xvals[,i], p[[i]][,3],
                     col=hpd.col[i], lty=hpd.lty[i], lwd=hpd.lwd[i])
        points(xvals[,i], p[[i]][,2], pch=point.pch[i],
               cex=point.cex[i], col=point.col[i], lwd=point.lwd[i])

    }


    if (show.legend)
        legend(legend.x, legend.y, gamma.names, col=point.col,
               pt.cex=point.cex, pch=point.pch, pt.lwd=point.lwd,
               horiz=legend.horiz,
               box.lwd=legend.box.lwd, cex=legend.cex,
               text.width=legend.width, xpd=legend.xpd, bty=legend.bty)
    axis(2)
    axis(1, at=c(-.5, 0, 1, 1.5), labels=c("", axis.labels, ""),
         tick=F)
        
    if (return.probs)
       p <- predprob.plist.fit(post, gamma, y, return.summary=FALSE)
}
